import random
import csv

random.seed(42)  # 固定种子确保可复现
total_groups = 100  # 目标总组数

# 定义参数选项（包含新增参数）
data_type = ["float16", "bfloat16"]
tokens = [128, 256, 512, 1024]
block_size_list = [128, 256]
hidden_size_list = [7168]  # 仅7168
cache_mode_list = [0]  # 仅0
head_num_list = [16, 32, 64, 128]
block_num = [192]  # 新增参数
hidden_size_wdqkv = [2112]  # 新增参数
hidden_size_wuq_head = [192]  # 新增参数
hidden_size_wuk_head = [512]  # 新增参数
hidden_size_wdq = [1536]  # 新增参数
hidden_size_rope_q_head = [64]  # 新增参数
hidden_size_rope_k = [64]  # 新增参数

rows = []
count = 0

while count < total_groups:
    # 随机选择参数（包括新增参数）
    token_num = random.choice(tokens)
    block_size = random.choice(block_size_list)
    hidden_size = random.choice(hidden_size_list)
    cache_mode = random.choice(cache_mode_list)
    dtype = random.choice(data_type)
    head_num = random.choice(head_num_list)
    b_num = random.choice(block_num)
    hs_wdqkv = random.choice(hidden_size_wdqkv)
    hs_wuq_head = random.choice(hidden_size_wuq_head)
    hs_wuk_head = random.choice(hidden_size_wuk_head)
    hs_wdq = random.choice(hidden_size_wdq)
    hs_rope_q_head = random.choice(hidden_size_rope_q_head)
    hs_rope_k = random.choice(hidden_size_rope_k)
    
    # 默认开启rmsNormQuant操作
    rms_norm_quant = True
    
    # 验证参数有效性
    valid = True
    
    # 约束1: tokenNum不能超过1024
    if token_num > 1024:
        valid = False
    
    # 约束2: blockSize必须是128或256
    if block_size not in [128, 256]:
        valid = False
    
    # 约束3: cacheMode为2或3时，blockSize必须为128（当前cache_mode仅0，仍保留逻辑）
    if cache_mode in [2, 3] and block_size != 128:
        valid = False
    
    # 约束4: hiddenSize必须在指定列表中
    if hidden_size not in hidden_size_list:
        valid = False
    
    # 约束5: 开启rmsNormQuant时的逻辑保持
    if not rms_norm_quant:
        if hidden_size != 6144 or cache_mode not in [0, 1]:
            valid = False
    
    if valid:
        rows.append({
            "tokenNum": token_num,
            "blockSize": block_size,
            "cacheMode": cache_mode,
            "hiddenSize": hidden_size,
            "dataType": dtype,
            "headNum": head_num,
            "blockNum": b_num,
            "hiddenSizeWDQKV": hs_wdqkv,
            "hiddenSizeWUQHead": hs_wuq_head,
            "hiddenSizeWUKHead": hs_wuk_head,
            "hiddenSizeWDQ": hs_wdq,
            "hiddenSizeRopeQHead": hs_rope_q_head,
            "hiddenSizeRopeK": hs_rope_k
        })
        count += 1

# 写入CSV文件（包含所有参数）
with open("./params/mlapo_data.csv", "w", newline="") as f:
    fieldnames = [
        "tokenNum", "blockSize", "cacheMode", "hiddenSize", 
        "dataType", "headNum", "blockNum", 
        "hiddenSizeWDQKV", "hiddenSizeWUQHead", "hiddenSizeWUKHead",
        "hiddenSizeWDQ", "hiddenSizeRopeQHead", "hiddenSizeRopeK"
    ]
    writer = csv.DictWriter(f, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(rows)

print(f"已生成{total_groups}组包含更新参数的测试数据，保存至mlapo_data.csv")